package org.example.SolutionNew003;

import sun.reflect.generics.tree.Tree;

import java.util.HashMap;
import java.util.Map;
import java.util.Stack;

public class DoubleTree009 {
    public static void main(String[] args) {
        DoubleTree009 dt = new DoubleTree009();
        TreeNode root = new TreeNode(3);
        root.left = new TreeNode(1);
        root.right = new TreeNode(4);
//        root.right.right = new TreeNode(7);
        root.left.right = new TreeNode(2);
        System.out.println(dt.kthSmallest(root, 4));
    }

    //使用栈
    public int kthSmallest1(TreeNode root, int k) {
        //中序遍历
        Stack<TreeNode> stack = new Stack<TreeNode>();
        int count = 0;

        while(!stack.isEmpty()||root != null) {
            while(root != null) {
                stack.push(root);
                root = root.left;
            }
            root = stack.pop();
            count++;
            if(count == k) {
                break;
            }
            root = root.right;
        }
        return root.val;
    }

    //遍历
    public int kthSmallest2(TreeNode root, int k) {
        return kthNode(root,k).val;
    }

    int count=0;
    public TreeNode kthNode(TreeNode root, int k) {
        if(root == null) return null;
        TreeNode left = kthNode(root.left,k);
        if(left != null)  return left;
        if(count==k-1){
            return root;
        }
        count++;
        TreeNode right = kthNode(root.right,k);
        if(right != null) return right;
        return null;
    }

    //记录节点个数
    public int kthSmallest(TreeNode root, int k) {
        MyBst bst=new MyBst(root,k);
        return bst.getKthNum();
    }
}

class MyBst{
    TreeNode root;
    Map<TreeNode,Integer> map;
    int k;

    public MyBst(TreeNode root,int k){
        this.root=root;
        this.k=k;
        map=new HashMap<TreeNode,Integer>();
        countNode(root);
    }

    public int getKthNum() {
        while(root!=null) {
            int count=map.getOrDefault(root.left,0);
            if(count<k-1){
                root=root.right;
                k-=count+1;
            }else if(count==k-1){
                return root.val;
            }else{
                root=root.left;
            }
        }
        return -1;
    }

    public int countNode(TreeNode node){
        if(node == null) return 0;

        map.put(node,1+countNode(node.left)+countNode(node.right));

        return map.get(node);
    }
}
